from langchain_community.vectorstores import Qdrant
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.chat_models import ChatZhipuAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain.docstore.document import Document
import gradio as gr
import os
import warnings
from tenacity import retry, stop_after_attempt, wait_exponential

# 禁用警告
warnings.filterwarnings('ignore')

# ---------- 直接设置API密钥 ----------
ZHIPUAI_API_KEY = "2f39319bdd864fc4a41bf6b8eed6efbc.uIsAkRrMwejVTIyc"  # 直接替换此处密钥

# ---------- 组件初始化（带重试机制） ----------
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def initialize_components():
    # 文本嵌入模型（显式指定API端点）
    embeddings = ZhipuAIEmbeddings(
        api_key=ZHIPUAI_API_KEY,
        model="glm-4",
        endpoint="https://open.bigmodel.cn/api/paas/v3/embeddings"
    )

    # 初始化向量数据库（必须包含至少一个文档）
    vectorstore = Qdrant.from_documents(
        documents=[Document(page_content="系统初始化文档")],
        embedding=embeddings,
        location=":memory:",
        collection_name="my_documents"
    )

    # 大语言模型配置
    llm = ChatZhipuAI(
        temperature=0.5,
        api_key=ZHIPUAI_API_KEY,
        model="glm-4",
        endpoint="https://open.bigmodel.cn/api/paas/v3/chat/completions"
    )

    # 文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )

    return embeddings, vectorstore, llm, text_splitter

# ...（保持其他函数不变）...

if __name__ == "__main__":
    # 调试输出
    print(f"[DEBUG] API密钥已直接设置，前6位：{ZHIPUAI_API_KEY[:6]}****")

    # 初始化组件
    embeddings, vectorstore, model_zhipu, text_splitter = initialize_components()
    retriever = vectorstore.as_retriever()

    # 启动服务
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )